import socket
import json
import logging
from dataclasses import dataclass
from typing import Dict, Any
from config import config

# Configure logging using settings from config
logging.basicConfig(
    level=getattr(logging, config.log_level),
    format=config.log_format
)
logger = logging.getLogger("unity-mcp-server")


@dataclass
class UnityConnection:
    """Manages the socket connection to the Unity Editor using UDP."""
    host: str = config.unity_host
    port: int = config.unity_port
    sock: socket.socket = None  # Socket for Unity communication

    def connect(self) -> bool:
        """Establish a UDP socket connection."""
        if self.sock:
            return True
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            logger.info(f"Created UDP socket for {self.host}:{self.port}")
            return True
        except Exception as e:
            logger.error(f"Failed to create UDP socket: {str(e)}")
            self.sock = None
            return False

    def disconnect(self):
        """Close the UDP connection."""
        if self.sock:
            try:
                self.sock.close()
            except Exception as e:
                logger.error(f"Error closing UDP socket: {str(e)}")
            finally:
                self.sock = None

    def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes:
        """Receive a complete UDP datagram response."""
        sock.settimeout(config.connection_timeout)
        try:
            data, addr = sock.recvfrom(buffer_size)
            logger.info(f"Received {len(data)} bytes from {addr}")

            decoded_data = data.decode('utf-8')

            # Ping response handling
            if decoded_data.strip().startswith('{"status":"success","result":{"message":"pong"'):
                logger.debug("Received ping response")
                return data

            # Content field processing
            if '"content":' in decoded_data:
                content_start = decoded_data.find('"content":') + 9
                content_end = decoded_data.rfind('"', content_start)
                if content_end > content_start:
                    content = decoded_data[content_start:content_end]
                    content = content.replace('\\"', '"')
                    decoded_data = decoded_data[:content_start] + \
                        content + decoded_data[content_end:]

            # Validate JSON
            json.loads(decoded_data)
            return data
        except json.JSONDecodeError as je:
            logger.error(f"JSON decode error: {str(je)}")
            raise Exception(f"Invalid JSON response: {str(je)}")
        except socket.timeout:
            logger.warning("Socket timeout during receive")
            raise Exception("Timeout receiving response")
        except Exception as e:
            logger.error(f"Error during receive: {str(e)}")
            raise

    def send_command(self, command_type: str, params: Dict[str, Any] = None) -> Dict[str, Any]:
        """Send a UDP datagram command to Unity."""
        if not self.sock and not self.connect():
            raise ConnectionError("Not connected to Unity")

        # Ping command handling
        if command_type == "ping":
            try:
                logger.debug("Sending UDP ping")
                self.sock.sendto(b"ping", (self.host, self.port))
                response_data = self.receive_full_response(self.sock)
                response = json.loads(response_data.decode('utf-8'))

                if response.get("status") != "success":
                    logger.warning("Ping response unsuccessful")
                    self.sock = None
                    raise ConnectionError("Connection verification failed")

                return {"message": "pong"}
            except Exception as e:
                logger.error(f"Ping error: {str(e)}")
                self.sock = None
                raise ConnectionError(
                    f"Connection verification failed: {str(e)}")

        # Normal command handling
        command = {"type": command_type, "params": params or {}}
        try:
            command_json = json.dumps(command, ensure_ascii=False)
            logger.info(f"Sending UDP command: {command_type}")

            self.sock.sendto(command_json.encode(
                'utf-8'), (self.host, self.port))
            response_data = self.receive_full_response(self.sock)

            response = json.loads(response_data.decode('utf-8'))
            if response.get("status") == "error":
                error_message = response.get("error") or "Unknown Unity error"
                logger.error(f"Unity error: {error_message}")
                raise Exception(error_message)

            return response.get("result", {})
        except Exception as e:
            logger.error(f"Communication error: {str(e)}")
            self.sock = None
            raise Exception(f"Failed to communicate: {str(e)}")


# Global Unity connection (remaining unchanged)
_unity_connection = None


def get_unity_connection() -> UnityConnection:
    """Retrieve or establish a persistent Unity connection."""
    global _unity_connection
    if _unity_connection is not None:
        try:
            _unity_connection.send_command("ping")
            logger.debug("Reusing existing UDP connection")
            return _unity_connection
        except Exception as e:
            logger.warning(f"Existing connection failed: {str(e)}")
            _unity_connection = None

    logger.info("Creating new UDP Unity connection")
    _unity_connection = UnityConnection()
    if not _unity_connection.connect():
        _unity_connection = None
        raise ConnectionError("Could not create UDP socket")

    try:
        _unity_connection.send_command("ping")
        logger.info("Verified UDP connection")
        return _unity_connection
    except Exception as e:
        logger.error(f"Connection verification failed: {str(e)}")
        _unity_connection = None
        raise ConnectionError(f"Could not establish UDP connection: {str(e)}")
